package cn.com.main;

import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

import java.io.IOException;

/**
 * 修改字节码
 *
 * @author LiuGuodong
 * @date 2024/03/04
 */
public class MonitorTransformer implements Opcodes {

    public static byte[] transform(byte[] classBytes) throws IOException {
        ClassReader reader = new ClassReader(classBytes);
        ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
        ClassVisitor visitor = new ClassVisitor(Opcodes.ASM5, writer) {
            @Override
            public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
                MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
                // 只为指定方法添加字节码插桩
                if ("method".equals(name) && "()V".equals(desc)) {
                    mv = new MethodVisitor(Opcodes.ASM5, mv) {
                        @Override
                        public void visitCode() {
                            super.visitCode();
                            // 在方法执行之前插入字节码
                            mv.visitMethodInsn(INVOKESTATIC, "TimerCountUtil", "start", "()V", false);
                        }

                        @Override
                        public void visitInsn(int opcode) {
                            // 在方法返回之前插入字节码
                            if (opcode == RETURN) {
                                mv.visitMethodInsn(INVOKESTATIC, "TimerCountUtil", "stop", "()V", false);
                            }
                            super.visitInsn(opcode);
                        }
                    };
                }
                return mv;
            }
        };
        reader.accept(visitor, ClassReader.EXPAND_FRAMES);
        return writer.toByteArray();
    }
}
